"""数据库连接管理

提供数据库连接、会话管理和连接池配置。
"""

from contextlib import contextmanager
from typing import Generator

from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.orm import sessionmaker, Session

from ..config.settings import get_database_url


class DatabaseManager:
    """数据库管理器"""

    def __init__(self):
        """初始化数据库管理器"""
        self._engine: Engine = None
        self._session_factory: sessionmaker = None

    @property
    def engine(self) -> Engine:
        """获取数据库引擎"""
        if self._engine is None:
            self._engine = create_engine(
                get_database_url(),
                pool_size=10,
                max_overflow=20,
                pool_pre_ping=True,
                pool_recycle=3600,
                echo=False
            )
        return self._engine

    @property
    def session_factory(self) -> sessionmaker:
        """获取会话工厂"""
        if self._session_factory is None:
            self._session_factory = sessionmaker(
                bind=self.engine,
                autocommit=False,
                autoflush=False
            )
        return self._session_factory

    def create_session(self) -> Session:
        """创建数据库会话
        
        Returns:
            数据库会话实例
        """
        return self.session_factory()

    @contextmanager
    def get_session(self) -> Generator[Session, None, None]:
        """获取数据库会话上下文管理器
        
        Yields:
            数据库会话实例
        """
        session = self.create_session()
        try:
            yield session
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    def test_connection(self) -> bool:
        """测试数据库连接
        
        Returns:
            连接是否成功
        """
        try:
            with self.get_session() as session:
                session.execute(text("SELECT 1"))
            return True
        except SQLAlchemyError:
            return False

    def close(self):
        """关闭数据库连接"""
        if self._engine:
            self._engine.dispose()
            self._engine = None
            self._session_factory = None


# 全局数据库管理器实例
db_manager = DatabaseManager()


def get_db_session() -> Generator[Session, None, None]:
    """获取数据库会话的依赖注入函数
    
    Yields:
        数据库会话实例
    """
    with db_manager.get_session() as session:
        yield session


def get_engine() -> Engine:
    """获取数据库引擎
    
    Returns:
        数据库引擎实例
    """
    return db_manager.engine


def test_database_connection() -> bool:
    """测试数据库连接
    
    Returns:
        连接是否成功
    """
    return db_manager.test_connection()
